import os
import random
import shutil

def merge_folders(folder_paths, target_folder, proportions):
    # 确保目标文件夹存在
    annotations_target = os.path.join(target_folder, "Annotations")
    images_target = os.path.join(target_folder, "JPEGImages")
    os.makedirs(annotations_target, exist_ok=True)
    os.makedirs(images_target, exist_ok=True)

    # 收集每个文件夹中的文件路径
    all_annotations = []
    all_images = []

    for folder_path in folder_paths:
        annotations_folder = os.path.join(folder_path, "Annotations")
        images_folder = os.path.join(folder_path, "JPEGImages")

        annotations = [os.path.join(annotations_folder, f) for f in os.listdir(annotations_folder) if f.endswith(".xml")]
        images = [os.path.join(images_folder, f) for f in os.listdir(images_folder) if f.endswith(('.jpg', '.png', '.jpeg'))]

        # 确保标签和图片数量一致
        assert len(annotations) == len(images), f"Mismatch between Annotations and JPEGImages in {folder_path}"

        all_annotations.append(annotations)
        all_images.append(images)

    # 按比例随机抽取
    for i, (annotations, images, proportion) in enumerate(zip(all_annotations, all_images, proportions)):
        if proportion == 1.0:
            selected_indices = range(len(annotations))  # 选择所有文件
        else:
            sample_size = int(len(annotations) * proportion)
            selected_indices = random.sample(range(len(annotations)), sample_size)

        for idx in selected_indices:
            annotation_src = annotations[idx]
            image_src = images[idx]

            annotation_dest = os.path.join(annotations_target, os.path.basename(annotation_src))
            image_dest = os.path.join(images_target, os.path.basename(image_src))

            shutil.copy2(annotation_src, annotation_dest)
            shutil.copy2(image_src, image_dest)

    print(f"Merge complete. Files saved to {target_folder}")

# 示例用法
folder_paths = ["folder1", "folder2", "folder3"]  # 替换为你的文件夹路径
target_folder = "merged_dataset"  # 目标文件夹路径
proportions = [0.5, 0.3, 0.2]  # 每个文件夹的抽取比例

merge_folders(folder_paths, target_folder, proportions)
